Code breaks on an oscillator-like problem

[evgenykurbatov]

Hello,

I have a system of two coupled ODEs. I tried several solvers from Diffrax, both non-stiff and stiff, with and without time step control. Every solver falls on this problem with the same error:

equinox.EquinoxRuntimeError: The maximum number of solver steps was reached. Try increasing `max_steps`.

However, they do solve the system when setting max_steps=1_000_000. The funny thing is that this is a test calculation to find a solution in a small vicinity of a center point of the ODE system, i.e. I expected the small amplitude oscillations around the initial values.

I suspect that I'm doing something wrong with Diffrax, because the scopy.integrate.solve_ivp does my system with DOP853 method. So the jax.experimental.ode.odeint does it too.

Here, the snippet:

import jax.numpy as jnp
import diffrax

def rhs(t, w, a):
    x, y = w
    xdot = - x**2 + y**2 - a**2
    ydot = - 2*x*y
    return xdot, ydot

a = 16.
w0 = (0., 1.1*a)

tbeg, tend, dt = 0., 1e3, 1.
saveat = diffrax.SaveAt(ts=jnp.arange(tbeg, tend, dt))

term = diffrax.ODETerm(rhs)
solver = diffrax.Dopri8()
stepsize_controller = diffrax.PIDController(rtol=1e-8, atol=1e-8)
#stepsize_controller = diffrax.PIDController(rtol=1e-3, atol=1e-6,
#                                            pcoeff=0.4, icoeff=0.3, dcoeff=0)
sol = diffrax.diffeqsolve(term, solver, tbeg, tend, 1e-3*dt, w0, args=a,
                          saveat=saveat, stepsize_controller=stepsize_controller)

[johannahaffner]

For me this terminates with a small number of steps when using the same tolerances as in solve_ivp. I played around with this and got the following:

• Using scipy.integrate.solve_ivp with their Dopri solver solves your system with 244910 evaluations of the vector field using their default tolerances (rtol=1e-3 and atol=1e-6).
• Using these same tolerances with the diffrax.PIDController solves the differential equation in only 12703 steps. The smaller number of steps taken at these tolerances is likely due to using a PID controller for the step size on the Diffrax side of things.
• When I decrease the tolerances on the Scipy side to match the rtol=atol=1e-8 you are using here, then Scipy takes 930326 steps.

There are a few things that make this comparison a little more tricky than meets the eye.

• Scipy does not seem to throw an error for very long solves with a large number of vector field evaluations, if I understand their documentation correctly (not a user myself). It does not even seem to have a max_steps argument (they have a max_step which refers to step size, not number of vector field evaluations).
• By contrast, Diffrax does throw an error if max_steps has been reached, allowing you to set a computational budget for your solve.

So part of you seeing this error, I think, just comes down to different behaviours for long integrations - Diffrax will let you know, Scipy will just keep going. In other words, the libraries have different concepts of what may constitute an "error", or a situation the user should be alerted about.

Decreasing the tolerances further (towards rtol=atol=1e-8), one seems to come to a point where the Diffrax solution will take more steps than the Scipy solution. I believe the likeliest explanation for this is that with very low tolerances, the step size controller will take very small steps for an oscillatory system. Understanding exactly how step size control is achieved using either solver will probably require taking a look at the Fortran source code of the Scipy solver, so that is tricky.
Comparisons between different solvers with differing implementation details are thus frequently done e.g. with work-precision diagrams. There is actually an addition to our FAQ in the current dev branch that tackles these questions, here.

Hope this helps!

[evgenykurbatov]

Oh, I got it. That was silly. Thank you!